{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PjRuzfwyImeC"
   },
   "source": [
    "# Adding Tools to NeMo Agent Toolkit\n",
    "\n",
    "In this notebook, we showcase how the NVIDIA NeMo Agent Toolkit (NAT) allows developers to easily add tool-calling capabilities to agents that integrate with the library. Tool calling allows chain-of-thought planners to augment the pretrained capabilities of LLMs with predefined methods and access to context retrieval services. This is a powerful paradigm in agentic AI development that has enabled revolutionary technologies like deep research and API-integrated AI systems.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Table of Contents\n",
    "\n",
    "- [0.0) Setup](#setup)\n",
    "  - [0.1) Prerequisites](#prereqs)\n",
    "  - [0.2) API Keys](#api-keys)\n",
    "  - [0.3) Data Sources](#data-sources)\n",
    "  - [0.4) Installing NeMo Agent Toolkit](#installing-nat)\n",
    "- [1.0) Creating a New Workflow](#creating-workflow)\n",
    "- [2.0) Defining New Tools](#defining-tools)\n",
    "  - [2.1) Total Product Sales Data Tool](#total-product-sales-tool)\n",
    "  - [2.2) Sales Per Day Tool](#sales-per-day-tool)\n",
    "  - [2.3) Detect Outliers Tool](#detect-outliers-tool)\n",
    "  - [2.4) Registering Tools](#registering-tools)\n",
    "  - [2.5) Updating The Configuration File](#update-config)\n",
    "  - [2.6) Running the Initial Workflow](#running-initial-workflow)\n",
    "- [3.0) Adding a Retrieval Tool using LlamaIndex](#adding-retrieval-tool)\n",
    "  - [3.1) Defining the Retrieval Tool](#defining-retrieval-tool)\n",
    "  - [3.2) Retrieval Tool Workflow Configuration File](#retrieval-config)\n",
    "  - [3.3) Running the Workflow](#running-retrieval-workflow)\n",
    "- [4.0) Next Steps](#next-steps)\n",
    "\n",
    "<span style=\"color:rgb(0, 31, 153); font-style: italic;\">Note: In Google Colab use the Table of Contents tab to navigate.</span>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"setup\"></a>\n",
    "# 0.0) Setup\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"prereqs\"></a>\n",
    "## 0.1) Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cwu2DOXB-u8M"
   },
   "source": [
    "- **Platform:** Linux, macOS, or Windows\n",
    "- **Python:** version 3.11, 3.12, or 3.13\n",
    "- **Python Packages:** `pip`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PzjU1lTaE3gW"
   },
   "source": [
    "<a id=\"api-keys\"></a>\n",
    "## 0.2) API Keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i4OTdB6wTdRZ"
   },
   "source": [
    "For this notebook, you will need the following API keys to run all examples end-to-end:\n",
    "\n",
    "- **NVIDIA Build:** You can obtain an NVIDIA Build API Key by creating an [NVIDIA Build](https://build.nvidia.com) account and generating a key at https://build.nvidia.com/settings/api-keys\n",
    "\n",
    "Then you can run the cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GBMnVYQ7E75x"
   },
   "source": [
    "<a id=\"data-sources\"></a>\n",
    "## 0.3) Data Sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ELmZ_Pdz-qX7"
   },
   "source": [
    "Several data files are required for this example. To keep this as a stand-alone example, the files are included here as cells which can be run to create them.\n",
    "\n",
    "The following cell creates the `data` directory as well as a `rag` subdirectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data/rag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e77jahmr_vdE"
   },
   "source": [
    "The following cell writes the `data/retail_sales_data.csv` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile data/retail_sales_data.csv\n",
    "Date,StoreID,Product,UnitsSold,Revenue,Promotion\n",
    "2024-01-01,S001,Laptop,1,1000,No\n",
    "2024-01-01,S001,Phone,9,4500,No\n",
    "2024-01-01,S001,Tablet,2,600,No\n",
    "2024-01-01,S002,Laptop,9,9000,No\n",
    "2024-01-01,S002,Phone,10,5000,No\n",
    "2024-01-01,S002,Tablet,5,1500,No\n",
    "2024-01-02,S001,Laptop,4,4000,No\n",
    "2024-01-02,S001,Phone,11,5500,No\n",
    "2024-01-02,S001,Tablet,7,2100,No\n",
    "2024-01-02,S002,Laptop,7,7000,No\n",
    "2024-01-02,S002,Phone,6,3000,No\n",
    "2024-01-02,S002,Tablet,9,2700,No\n",
    "2024-01-03,S001,Laptop,6,6000,No\n",
    "2024-01-03,S001,Phone,7,3500,No\n",
    "2024-01-03,S001,Tablet,8,2400,No\n",
    "2024-01-03,S002,Laptop,3,3000,No\n",
    "2024-01-03,S002,Phone,16,8000,No\n",
    "2024-01-03,S002,Tablet,5,1500,No\n",
    "2024-01-04,S001,Laptop,5,5000,No\n",
    "2024-01-04,S001,Phone,11,5500,No\n",
    "2024-01-04,S001,Tablet,9,2700,No\n",
    "2024-01-04,S002,Laptop,2,2000,No\n",
    "2024-01-04,S002,Phone,12,6000,No\n",
    "2024-01-04,S002,Tablet,7,2100,No\n",
    "2024-01-05,S001,Laptop,8,8000,No\n",
    "2024-01-05,S001,Phone,18,9000,No\n",
    "2024-01-05,S001,Tablet,5,1500,No\n",
    "2024-01-05,S002,Laptop,7,7000,No\n",
    "2024-01-05,S002,Phone,10,5000,No\n",
    "2024-01-05,S002,Tablet,10,3000,No\n",
    "2024-01-06,S001,Laptop,9,9000,No\n",
    "2024-01-06,S001,Phone,11,5500,No\n",
    "2024-01-06,S001,Tablet,5,1500,No\n",
    "2024-01-06,S002,Laptop,5,5000,No\n",
    "2024-01-06,S002,Phone,14,7000,No\n",
    "2024-01-06,S002,Tablet,10,3000,No\n",
    "2024-01-07,S001,Laptop,2,2000,No\n",
    "2024-01-07,S001,Phone,15,7500,No\n",
    "2024-01-07,S001,Tablet,6,1800,No\n",
    "2024-01-07,S002,Laptop,0,0,No\n",
    "2024-01-07,S002,Phone,7,3500,No\n",
    "2024-01-07,S002,Tablet,12,3600,No\n",
    "2024-01-08,S001,Laptop,5,5000,No\n",
    "2024-01-08,S001,Phone,8,4000,No\n",
    "2024-01-08,S001,Tablet,5,1500,No\n",
    "2024-01-08,S002,Laptop,4,4000,No\n",
    "2024-01-08,S002,Phone,11,5500,No\n",
    "2024-01-08,S002,Tablet,9,2700,No\n",
    "2024-01-09,S001,Laptop,6,6000,No\n",
    "2024-01-09,S001,Phone,9,4500,No\n",
    "2024-01-09,S001,Tablet,8,2400,No\n",
    "2024-01-09,S002,Laptop,7,7000,No\n",
    "2024-01-09,S002,Phone,11,5500,No\n",
    "2024-01-09,S002,Tablet,8,2400,No\n",
    "2024-01-10,S001,Laptop,6,6000,No\n",
    "2024-01-10,S001,Phone,11,5500,No\n",
    "2024-01-10,S001,Tablet,5,1500,No\n",
    "2024-01-10,S002,Laptop,8,8000,No\n",
    "2024-01-10,S002,Phone,5,2500,No\n",
    "2024-01-10,S002,Tablet,6,1800,No\n",
    "2024-01-11,S001,Laptop,5,5000,No\n",
    "2024-01-11,S001,Phone,7,3500,No\n",
    "2024-01-11,S001,Tablet,5,1500,No\n",
    "2024-01-11,S002,Laptop,4,4000,No\n",
    "2024-01-11,S002,Phone,10,5000,No\n",
    "2024-01-11,S002,Tablet,4,1200,No\n",
    "2024-01-12,S001,Laptop,2,2000,No\n",
    "2024-01-12,S001,Phone,10,5000,No\n",
    "2024-01-12,S001,Tablet,9,2700,No\n",
    "2024-01-12,S002,Laptop,8,8000,No\n",
    "2024-01-12,S002,Phone,10,5000,No\n",
    "2024-01-12,S002,Tablet,14,4200,No\n",
    "2024-01-13,S001,Laptop,3,3000,No\n",
    "2024-01-13,S001,Phone,6,3000,No\n",
    "2024-01-13,S001,Tablet,9,2700,No\n",
    "2024-01-13,S002,Laptop,1,1000,No\n",
    "2024-01-13,S002,Phone,12,6000,No\n",
    "2024-01-13,S002,Tablet,7,2100,No\n",
    "2024-01-14,S001,Laptop,4,4000,Yes\n",
    "2024-01-14,S001,Phone,16,8000,Yes\n",
    "2024-01-14,S001,Tablet,4,1200,Yes\n",
    "2024-01-14,S002,Laptop,5,5000,Yes\n",
    "2024-01-14,S002,Phone,14,7000,Yes\n",
    "2024-01-14,S002,Tablet,6,1800,Yes\n",
    "2024-01-15,S001,Laptop,9,9000,No\n",
    "2024-01-15,S001,Phone,6,3000,No\n",
    "2024-01-15,S001,Tablet,11,3300,No\n",
    "2024-01-15,S002,Laptop,5,5000,No\n",
    "2024-01-15,S002,Phone,10,5000,No\n",
    "2024-01-15,S002,Tablet,4,1200,No\n",
    "2024-01-16,S001,Laptop,6,6000,No\n",
    "2024-01-16,S001,Phone,11,5500,No\n",
    "2024-01-16,S001,Tablet,5,1500,No\n",
    "2024-01-16,S002,Laptop,4,4000,No\n",
    "2024-01-16,S002,Phone,7,3500,No\n",
    "2024-01-16,S002,Tablet,4,1200,No\n",
    "2024-01-17,S001,Laptop,6,6000,No\n",
    "2024-01-17,S001,Phone,14,7000,No\n",
    "2024-01-17,S001,Tablet,7,2100,No\n",
    "2024-01-17,S002,Laptop,3,3000,No\n",
    "2024-01-17,S002,Phone,7,3500,No\n",
    "2024-01-17,S002,Tablet,6,1800,No\n",
    "2024-01-18,S001,Laptop,7,7000,Yes\n",
    "2024-01-18,S001,Phone,10,5000,Yes\n",
    "2024-01-18,S001,Tablet,6,1800,Yes\n",
    "2024-01-18,S002,Laptop,5,5000,Yes\n",
    "2024-01-18,S002,Phone,16,8000,Yes\n",
    "2024-01-18,S002,Tablet,8,2400,Yes\n",
    "2024-01-19,S001,Laptop,4,4000,No\n",
    "2024-01-19,S001,Phone,12,6000,No\n",
    "2024-01-19,S001,Tablet,7,2100,No\n",
    "2024-01-19,S002,Laptop,3,3000,No\n",
    "2024-01-19,S002,Phone,12,6000,No\n",
    "2024-01-19,S002,Tablet,8,2400,No\n",
    "2024-01-20,S001,Laptop,6,6000,No\n",
    "2024-01-20,S001,Phone,8,4000,No\n",
    "2024-01-20,S001,Tablet,6,1800,No\n",
    "2024-01-20,S002,Laptop,8,8000,No\n",
    "2024-01-20,S002,Phone,9,4500,No\n",
    "2024-01-20,S002,Tablet,8,2400,No\n",
    "2024-01-21,S001,Laptop,3,3000,No\n",
    "2024-01-21,S001,Phone,9,4500,No\n",
    "2024-01-21,S001,Tablet,5,1500,No\n",
    "2024-01-21,S002,Laptop,8,8000,No\n",
    "2024-01-21,S002,Phone,15,7500,No\n",
    "2024-01-21,S002,Tablet,7,2100,No\n",
    "2024-01-22,S001,Laptop,1,1000,No\n",
    "2024-01-22,S001,Phone,15,7500,No\n",
    "2024-01-22,S001,Tablet,5,1500,No\n",
    "2024-01-22,S002,Laptop,11,11000,No\n",
    "2024-01-22,S002,Phone,4,2000,No\n",
    "2024-01-22,S002,Tablet,4,1200,No\n",
    "2024-01-23,S001,Laptop,3,3000,No\n",
    "2024-01-23,S001,Phone,8,4000,No\n",
    "2024-01-23,S001,Tablet,8,2400,No\n",
    "2024-01-23,S002,Laptop,6,6000,No\n",
    "2024-01-23,S002,Phone,12,6000,No\n",
    "2024-01-23,S002,Tablet,12,3600,No\n",
    "2024-01-24,S001,Laptop,2,2000,No\n",
    "2024-01-24,S001,Phone,14,7000,No\n",
    "2024-01-24,S001,Tablet,6,1800,No\n",
    "2024-01-24,S002,Laptop,1,1000,No\n",
    "2024-01-24,S002,Phone,5,2500,No\n",
    "2024-01-24,S002,Tablet,7,2100,No\n",
    "2024-01-25,S001,Laptop,7,7000,No\n",
    "2024-01-25,S001,Phone,11,5500,No\n",
    "2024-01-25,S001,Tablet,11,3300,No\n",
    "2024-01-25,S002,Laptop,6,6000,No\n",
    "2024-01-25,S002,Phone,11,5500,No\n",
    "2024-01-25,S002,Tablet,5,1500,No\n",
    "2024-01-26,S001,Laptop,5,5000,Yes\n",
    "2024-01-26,S001,Phone,22,11000,Yes\n",
    "2024-01-26,S001,Tablet,7,2100,Yes\n",
    "2024-01-26,S002,Laptop,6,6000,Yes\n",
    "2024-01-26,S002,Phone,24,12000,Yes\n",
    "2024-01-26,S002,Tablet,3,900,Yes\n",
    "2024-01-27,S001,Laptop,7,7000,Yes\n",
    "2024-01-27,S001,Phone,20,10000,Yes\n",
    "2024-01-27,S001,Tablet,6,1800,Yes\n",
    "2024-01-27,S002,Laptop,4,4000,Yes\n",
    "2024-01-27,S002,Phone,8,4000,Yes\n",
    "2024-01-27,S002,Tablet,6,1800,Yes\n",
    "2024-01-28,S001,Laptop,10,10000,No\n",
    "2024-01-28,S001,Phone,15,7500,No\n",
    "2024-01-28,S001,Tablet,12,3600,No\n",
    "2024-01-28,S002,Laptop,6,6000,No\n",
    "2024-01-28,S002,Phone,11,5500,No\n",
    "2024-01-28,S002,Tablet,10,3000,No\n",
    "2024-01-29,S001,Laptop,3,3000,No\n",
    "2024-01-29,S001,Phone,16,8000,No\n",
    "2024-01-29,S001,Tablet,5,1500,No\n",
    "2024-01-29,S002,Laptop,6,6000,No\n",
    "2024-01-29,S002,Phone,17,8500,No\n",
    "2024-01-29,S002,Tablet,2,600,No\n",
    "2024-01-30,S001,Laptop,3,3000,No\n",
    "2024-01-30,S001,Phone,11,5500,No\n",
    "2024-01-30,S001,Tablet,2,600,No\n",
    "2024-01-30,S002,Laptop,6,6000,No\n",
    "2024-01-30,S002,Phone,16,8000,No\n",
    "2024-01-30,S002,Tablet,8,2400,No\n",
    "2024-01-31,S001,Laptop,5,5000,Yes\n",
    "2024-01-31,S001,Phone,22,11000,Yes\n",
    "2024-01-31,S001,Tablet,9,2700,Yes\n",
    "2024-01-31,S002,Laptop,3,3000,Yes\n",
    "2024-01-31,S002,Phone,14,7000,Yes\n",
    "2024-01-31,S002,Tablet,4,1200,Yes\n",
    "2024-02-01,S001,Laptop,2,2000,No\n",
    "2024-02-01,S001,Phone,7,3500,No\n",
    "2024-02-01,S001,Tablet,11,3300,No\n",
    "2024-02-01,S002,Laptop,6,6000,No\n",
    "2024-02-01,S002,Phone,11,5500,No\n",
    "2024-02-01,S002,Tablet,5,1500,No\n",
    "2024-02-02,S001,Laptop,2,2000,No\n",
    "2024-02-02,S001,Phone,9,4500,No\n",
    "2024-02-02,S001,Tablet,7,2100,No\n",
    "2024-02-02,S002,Laptop,5,5000,No\n",
    "2024-02-02,S002,Phone,9,4500,No\n",
    "2024-02-02,S002,Tablet,12,3600,No\n",
    "2024-02-03,S001,Laptop,9,9000,No\n",
    "2024-02-03,S001,Phone,12,6000,No\n",
    "2024-02-03,S001,Tablet,9,2700,No\n",
    "2024-02-03,S002,Laptop,10,10000,No\n",
    "2024-02-03,S002,Phone,6,3000,No\n",
    "2024-02-03,S002,Tablet,10,3000,No\n",
    "2024-02-04,S001,Laptop,6,6000,No\n",
    "2024-02-04,S001,Phone,5,2500,No\n",
    "2024-02-04,S001,Tablet,8,2400,No\n",
    "2024-02-04,S002,Laptop,6,6000,No\n",
    "2024-02-04,S002,Phone,10,5000,No\n",
    "2024-02-04,S002,Tablet,10,3000,No\n",
    "2024-02-05,S001,Laptop,7,7000,No\n",
    "2024-02-05,S001,Phone,13,6500,No\n",
    "2024-02-05,S001,Tablet,11,3300,No\n",
    "2024-02-05,S002,Laptop,8,8000,No\n",
    "2024-02-05,S002,Phone,11,5500,No\n",
    "2024-02-05,S002,Tablet,8,2400,No\n",
    "2024-02-06,S001,Laptop,5,5000,No\n",
    "2024-02-06,S001,Phone,14,7000,No\n",
    "2024-02-06,S001,Tablet,4,1200,No\n",
    "2024-02-06,S002,Laptop,2,2000,No\n",
    "2024-02-06,S002,Phone,11,5500,No\n",
    "2024-02-06,S002,Tablet,7,2100,No\n",
    "2024-02-07,S001,Laptop,6,6000,No\n",
    "2024-02-07,S001,Phone,7,3500,No\n",
    "2024-02-07,S001,Tablet,9,2700,No\n",
    "2024-02-07,S002,Laptop,2,2000,No\n",
    "2024-02-07,S002,Phone,8,4000,No\n",
    "2024-02-07,S002,Tablet,9,2700,No\n",
    "2024-02-08,S001,Laptop,5,5000,No\n",
    "2024-02-08,S001,Phone,12,6000,No\n",
    "2024-02-08,S001,Tablet,3,900,No\n",
    "2024-02-08,S002,Laptop,8,8000,No\n",
    "2024-02-08,S002,Phone,5,2500,No\n",
    "2024-02-08,S002,Tablet,8,2400,No\n",
    "2024-02-09,S001,Laptop,6,6000,Yes\n",
    "2024-02-09,S001,Phone,18,9000,Yes\n",
    "2024-02-09,S001,Tablet,5,1500,Yes\n",
    "2024-02-09,S002,Laptop,7,7000,Yes\n",
    "2024-02-09,S002,Phone,18,9000,Yes\n",
    "2024-02-09,S002,Tablet,5,1500,Yes\n",
    "2024-02-10,S001,Laptop,9,9000,No\n",
    "2024-02-10,S001,Phone,6,3000,No\n",
    "2024-02-10,S001,Tablet,8,2400,No\n",
    "2024-02-10,S002,Laptop,7,7000,No\n",
    "2024-02-10,S002,Phone,5,2500,No\n",
    "2024-02-10,S002,Tablet,6,1800,No\n",
    "2024-02-11,S001,Laptop,6,6000,No\n",
    "2024-02-11,S001,Phone,11,5500,No\n",
    "2024-02-11,S001,Tablet,2,600,No\n",
    "2024-02-11,S002,Laptop,7,7000,No\n",
    "2024-02-11,S002,Phone,5,2500,No\n",
    "2024-02-11,S002,Tablet,9,2700,No\n",
    "2024-02-12,S001,Laptop,5,5000,No\n",
    "2024-02-12,S001,Phone,5,2500,No\n",
    "2024-02-12,S001,Tablet,4,1200,No\n",
    "2024-02-12,S002,Laptop,1,1000,No\n",
    "2024-02-12,S002,Phone,14,7000,No\n",
    "2024-02-12,S002,Tablet,15,4500,No\n",
    "2024-02-13,S001,Laptop,3,3000,No\n",
    "2024-02-13,S001,Phone,18,9000,No\n",
    "2024-02-13,S001,Tablet,8,2400,No\n",
    "2024-02-13,S002,Laptop,5,5000,No\n",
    "2024-02-13,S002,Phone,8,4000,No\n",
    "2024-02-13,S002,Tablet,6,1800,No\n",
    "2024-02-14,S001,Laptop,4,4000,No\n",
    "2024-02-14,S001,Phone,9,4500,No\n",
    "2024-02-14,S001,Tablet,6,1800,No\n",
    "2024-02-14,S002,Laptop,4,4000,No\n",
    "2024-02-14,S002,Phone,6,3000,No\n",
    "2024-02-14,S002,Tablet,7,2100,No\n",
    "2024-02-15,S001,Laptop,4,4000,Yes\n",
    "2024-02-15,S001,Phone,26,13000,Yes\n",
    "2024-02-15,S001,Tablet,5,1500,Yes\n",
    "2024-02-15,S002,Laptop,2,2000,Yes\n",
    "2024-02-15,S002,Phone,14,7000,Yes\n",
    "2024-02-15,S002,Tablet,6,1800,Yes\n",
    "2024-02-16,S001,Laptop,7,7000,No\n",
    "2024-02-16,S001,Phone,9,4500,No\n",
    "2024-02-16,S001,Tablet,1,300,No\n",
    "2024-02-16,S002,Laptop,6,6000,No\n",
    "2024-02-16,S002,Phone,12,6000,No\n",
    "2024-02-16,S002,Tablet,10,3000,No\n",
    "2024-02-17,S001,Laptop,5,5000,No\n",
    "2024-02-17,S001,Phone,8,4000,No\n",
    "2024-02-17,S001,Tablet,14,4200,No\n",
    "2024-02-17,S002,Laptop,4,4000,No\n",
    "2024-02-17,S002,Phone,13,6500,No\n",
    "2024-02-17,S002,Tablet,7,2100,No\n",
    "2024-02-18,S001,Laptop,6,6000,Yes\n",
    "2024-02-18,S001,Phone,22,11000,Yes\n",
    "2024-02-18,S001,Tablet,9,2700,Yes\n",
    "2024-02-18,S002,Laptop,2,2000,Yes\n",
    "2024-02-18,S002,Phone,10,5000,Yes\n",
    "2024-02-18,S002,Tablet,12,3600,Yes\n",
    "2024-02-19,S001,Laptop,6,6000,No\n",
    "2024-02-19,S001,Phone,12,6000,No\n",
    "2024-02-19,S001,Tablet,3,900,No\n",
    "2024-02-19,S002,Laptop,3,3000,No\n",
    "2024-02-19,S002,Phone,4,2000,No\n",
    "2024-02-19,S002,Tablet,7,2100,No\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RqbTkKoX_81n"
   },
   "source": [
    "The following cell writes the RAG product catalog file, `data/rag/product_catalog.md`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form"
   },
   "outputs": [],
   "source": [
    "%%writefile data/rag/product_catalog.md\n",
    "# Product Catalog: Smartphones, Laptops, and Tablets\n",
    "\n",
    "## Smartphones\n",
    "\n",
    "The Veltrix Solis Z9 is a flagship device in the premium smartphone segment. It builds on a decade of design iterations that prioritize screen-to-body ratio, minimal bezels, and high refresh rate displays. The 6.7-inch AMOLED panel with 120Hz refresh rate delivers immersive visual experiences, whether in gaming, video streaming, or augmented reality applications. The display's GorillaGlass Fusion coating provides scratch resistance and durability, and the thin form factor is engineered using a titanium-aluminum alloy chassis to reduce weight without compromising rigidity.\n",
    "\n",
    "Internally, the Solis Z9 is powered by the OrionEdge V14 chipset, a 4nm process SoC designed for high-efficiency workloads. Its AI accelerator module handles on-device tasks such as voice transcription, camera optimization, and intelligent background app management. The inclusion of 12GB LPDDR5 RAM and a 256GB UFS 3.1 storage system allows for seamless multitasking, instant app launching, and rapid data access. The device supports eSIM and dual physical SIM configurations, catering to global travelers and hybrid network users.\n",
    "\n",
    "Photography and videography are central to the Solis Z9 experience. The triple-camera system incorporates a periscope-style 8MP telephoto lens with 5x optical zoom, a 12MP ultra-wide sensor with macro capabilities, and a 64MP main sensor featuring optical image stabilization (OIS) and phase detection autofocus (PDAF). Night mode and HDRX+ processing enable high-fidelity image capture in challenging lighting conditions.\n",
    "\n",
    "Software-wise, the device ships with LunOS 15, a lightweight Android fork optimized for modular updates and privacy compliance. The system supports secure containers for work profiles and AI-powered notifications that summarize app alerts across channels. Facial unlock is augmented by a 3D IR depth sensor, providing reliable biometric security alongside the ultrasonic in-display fingerprint scanner.\n",
    "\n",
    "The Solis Z9 is a culmination of over a decade of design experimentation in mobile form factors, ranging from curved-edge screens to under-display camera arrays. Its balance of performance, battery efficiency, and user-centric software makes it an ideal daily driver for content creators, mobile gamers, and enterprise users.\n",
    "\n",
    "## Laptops\n",
    "\n",
    "The Cryon Vanta 16X represents the latest evolution of portable computing power tailored for professional-grade workloads.\n",
    "\n",
    "The Vanta 16X features a unibody chassis milled from aircraft-grade aluminum using CNC machining. The thermal design integrates vapor chamber cooling and dual-fan exhaust architecture to support sustained performance under high computational loads. The 16-inch 4K UHD display is color-calibrated at the factory and supports HDR10+, making it suitable for cinematic video editing and high-fidelity CAD modeling.\n",
    "\n",
    "Powering the device is Intel's Core i9-13900H processor, which includes 14 cores with a hybrid architecture combining performance and efficiency cores. This allows the system to dynamically balance power consumption and raw speed based on active workloads. The dedicated Zephira RTX 4700G GPU features 8GB of GDDR6 VRAM and is optimized for CUDA and Tensor Core operations, enabling applications in real-time ray tracing, AI inference, and 3D rendering.\n",
    "\n",
    "The Vanta 16X includes a 2TB PCIe Gen 4 NVMe SSD, delivering sequential read/write speeds above 7GB/s, and 32GB of high-bandwidth DDR5 RAM. The machine supports hardware-accelerated virtualization and dual-booting, and ships with VireoOS Pro pre-installed, with official drivers available for Fedora, Ubuntu LTS, and NebulaOS.\n",
    "\n",
    "Input options are expansive. The keyboard features per-key RGB lighting and programmable macros, while the haptic touchpad supports multi-gesture navigation and palm rejection. Port variety includes dual Thunderbolt 4 ports, a full-size SD Express card reader, HDMI 2.1, 2.5G Ethernet, three USB-A 3.2 ports, and a 3.5mm TRRS audio jack. A fingerprint reader is embedded in the power button and supports biometric logins via Windows Hello.\n",
    "\n",
    "The history of the Cryon laptop line dates back to the early 2010s, when the company launched its first ultrabook aimed at mobile developers. Since then, successive generations have introduced carbon fiber lids, modular SSD bays, and convertible form factors. The Vanta 16X continues this tradition by integrating a customizable BIOS, a modular fan assembly, and a trackpad optimized for creative software like Blender and Adobe Creative Suite.\n",
    "\n",
    "Designed for software engineers, data scientists, film editors, and 3D artists, the Cryon Vanta 16X is a workstation-class laptop in a portable shell.\n",
    "\n",
    "## Tablets\n",
    "\n",
    "The Nebulyn Ark S12 Ultra reflects the current apex of tablet technology, combining high-end hardware with software environments tailored for productivity and creativity.\n",
    "\n",
    "The Ark S12 Ultra is built around a 12.9-inch OLED display that supports 144Hz refresh rate and HDR10+ dynamic range. With a resolution of 2800 x 1752 pixels and a contrast ratio of 1,000,000:1, the screen delivers vibrant color reproduction ideal for design and media consumption. The display supports true tone adaptation and low blue-light filtering for prolonged use.\n",
    "\n",
    "Internally, the tablet uses Qualcomm's Snapdragon 8 Gen 3 SoC, which includes an Adreno 750 GPU and an NPU for on-device AI tasks. The device ships with 16GB LPDDR5X RAM and 512GB of storage with support for NVMe expansion via a proprietary magnetic dock. The 11200mAh battery enables up to 15 hours of typical use and recharges to 80 percent in 45 minutes via 45W USB-C PD.\n",
    "\n",
    "The Ark's history traces back to the original Nebulyn Tab, which launched in 2014 as an e-reader and video streaming device. Since then, the line has evolved through multiple iterations that introduced stylus support, high-refresh screens, and multi-window desktop modes. The current model supports NebulynVerse, a DeX-like environment that allows external display mirroring and full multitasking with overlapping windows and keyboard shortcuts.\n",
    "\n",
    "Input capabilities are central to the Ark S12 Ultra’s appeal. The Pluma Stylus 3 features magnetic charging, 4096 pressure levels, and tilt detection. It integrates haptic feedback to simulate traditional pen strokes and brush textures. The device also supports a SnapCover keyboard that includes a trackpad and programmable shortcut keys. With the stylus and keyboard, users can effectively transform the tablet into a mobile workstation or digital sketchbook.\n",
    "\n",
    "Camera hardware includes a 13MP main sensor and a 12MP ultra-wide front camera with center-stage tracking and biometric unlock. Microphone arrays with beamforming enable studio-quality call audio. Connectivity includes Wi-Fi 7, Bluetooth 5.3, and optional LTE/5G with eSIM.\n",
    "\n",
    "Software support is robust. The device runs NebulynOS 6.0, based on Android 14L, and supports app sandboxing, multi-user profiles, and remote device management. Integration with cloud services, including SketchNimbus and ThoughtSpace, allows for real-time collaboration and syncing of content across devices.\n",
    "\n",
    "This tablet is targeted at professionals who require a balance between media consumption, creativity, and light productivity. Typical users include architects, consultants, university students, and UX designers.\n",
    "\n",
    "## Comparative Summary\n",
    "\n",
    "Each of these devices—the Veltrix Solis Z9, Cryon Vanta 16X, and Nebulyn Ark S12 Ultra—represents a best-in-class interpretation of its category. The Solis Z9 excels in mobile photography and everyday communication. The Vanta 16X is tailored for high-performance applications such as video production and AI prototyping. The Ark S12 Ultra provides a canvas for creativity, note-taking, and hybrid productivity use cases.\n",
    "\n",
    "## Historical Trends and Design Evolution\n",
    "\n",
    "Design across all three categories is converging toward modularity, longevity, and environmental sustainability. Recycled materials, reparability scores, and software longevity are becoming integral to brand reputation and product longevity. Future iterations are expected to feature tighter integration with wearable devices, ambient AI experiences, and cross-device workflows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0IUUGtXSFB5G"
   },
   "source": [
    "<a id=\"installing-nat\"></a>\n",
    "## 0.4) Installing NeMo Agent Toolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OSICVNHGGm9l"
   },
   "source": [
    "The recommended way to install NAT is through `pip` or `uv pip`.\n",
    "\n",
    "First, we will install `uv` which offers parallel downloads and faster dependency resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install uv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EBV2Gh9NIC8R"
   },
   "source": [
    "NeMo Agent toolkit can be installed through the PyPI `nvidia-nat` package.\n",
    "\n",
    "There are several optional subpackages available for NAT. For this example, we will rely on two subpackages:\n",
    "* The `langchain` subpackage contains useful components for integrating and running within [LangChain](https://python.langchain.com/docs/introduction/).\n",
    "* The `llama-index` subpackage contains useful components for integrating and running within [LlamaIndex](https://developers.llamaindex.ai/python/framework/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "uv pip show -q \"nvidia-nat-langchain\"\n",
    "nat_langchain_installed=$?\n",
    "uv pip show -q \"nvidia-nat-llama-index\"\n",
    "nat_llama_index_installed=$?\n",
    "if [[ ${nat_langchain_installed} -ne 0 || ${nat_llama_index_installed} -ne 0 ]]; then\n",
    "    uv pip install \"nvidia-nat[langchain,llama-index]\"\n",
    "else\n",
    "    echo \"nvidia-nat[langchain,llama-index] is already installed\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l7kWJ8yeJJhQ"
   },
   "source": [
    "<a id=\"creating-workflow\"></a>\n",
    "# 1.0) Creating a New Workflow\n",
    "\n",
    "As explained in detail in previous notebooks in this series, we can use the `nat workflow create` sub-command to create the necessary directory structure for a new agent.\n",
    "\n",
    "Within this directory we can define all of the functions that we want to be available to the agent at runtime. In this notebook specifically we are going to demonstrate the integration of new tools to the workflow. We will make them available or 'discoverable' by the agent by defining these tool calls within the `register.py` function.\n",
    "\n",
    "In the next notebook of this series we will show how to accomplish multi-agent orchestration via this same approach, but for simplicity, we will keep this notebook focused on tool calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create retail_sales_agent_nb3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iSDMOrSQKtBr"
   },
   "source": [
    "<a id=\"defining-tools\"></a>\n",
    "# 2.0) Defining New Tools\n",
    "\n",
    "Next we will show you how to add new tools to the agent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PKABb9h0ej1z"
   },
   "source": [
    "<a id=\"total-product-sales-tool\"></a>\n",
    "## 2.1) Total Product Sales Data Tool\n",
    "\n",
    "This tool gets total sales for a specific product from `data/product_catalog.md`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb3/src/retail_sales_agent_nb3/total_product_sales_data_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class GetTotalProductSalesDataConfig(FunctionBaseConfig, name=\"get_total_product_sales_data\"):\n",
    "    \"\"\"Get total sales data by product.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=GetTotalProductSalesDataConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def get_total_product_sales_data_function(config: GetTotalProductSalesDataConfig, _builder: Builder):\n",
    "    \"\"\"Get total sales data for a specific product.\"\"\"\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _get_total_product_sales_data(product_name: str) -> str:\n",
    "        \"\"\"\n",
    "        Retrieve total sales data for a specific product.\n",
    "\n",
    "        Args:\n",
    "            product_name: Name of the product\n",
    "\n",
    "        Returns:\n",
    "            String message containing total sales data\n",
    "        \"\"\"\n",
    "        df['Product'] = df[\"Product\"].apply(lambda x: x.lower())\n",
    "        revenue = df[df['Product'] == product_name]['Revenue'].sum()\n",
    "        units_sold = df[df['Product'] == product_name]['UnitsSold'].sum()\n",
    "\n",
    "        return f\"Revenue for {product_name} are {revenue} and total units sold are {units_sold}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _get_total_product_sales_data,\n",
    "        description=_get_total_product_sales_data.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cg7cKMTPe26D"
   },
   "source": [
    "<a id=\"sales-per-day-tool\"></a>\n",
    "## 2.2) Sales Per Day Tool\n",
    "\n",
    "This tool gets the total sales across all products per day from `data/product_catalog.md`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb3/src/retail_sales_agent_nb3/sales_per_day_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class GetSalesPerDayConfig(FunctionBaseConfig, name=\"get_sales_per_day\"):\n",
    "    \"\"\"Get total sales across all products per day.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=GetSalesPerDayConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def sales_per_day_function(config: GetSalesPerDayConfig, builder: Builder):\n",
    "    \"\"\"Get total sales across all products per day.\"\"\"\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "    df['Product'] = df[\"Product\"].apply(lambda x: x.lower())\n",
    "\n",
    "    async def _get_sales_per_day(date: str, product: str) -> str:\n",
    "        \"\"\"\n",
    "        Calculate total sales data across all products for a specific date.\n",
    "\n",
    "        Args:\n",
    "            date: Date in YYYY-MM-DD format\n",
    "            product: Product name\n",
    "\n",
    "        Returns:\n",
    "            String message with the total sales for the day\n",
    "        \"\"\"\n",
    "        if date == \"None\":\n",
    "            return \"Please provide a date in YYYY-MM-DD format.\"\n",
    "        total_revenue = df[(df['Date'] == date) & (df['Product'] == product)]['Revenue'].sum()\n",
    "        total_units_sold = df[(df['Date'] == date) & (df['Product'] == product)]['UnitsSold'].sum()\n",
    "\n",
    "        return f\"Total revenue for {date} is {total_revenue} and total units sold is {total_units_sold}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _get_sales_per_day,\n",
    "        description=_get_sales_per_day.__doc__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OatXydqNe6gf"
   },
   "source": [
    "<a id=\"detect-outliers-tool\"></a>\n",
    "## 2.3) Detect Outliers Tool\n",
    "\n",
    "This tool detects outliers in `data/product_catalog.md` data using IQR (Interquartile Range) method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb3/src/retail_sales_agent_nb3/detect_outliers_tool.py\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "\n",
    "class DetectOutliersIQRConfig(FunctionBaseConfig, name=\"detect_outliers_iqr\"):\n",
    "    \"\"\"Detect outliers in sales data using IQR method.\"\"\"\n",
    "    data_path: str = Field(description=\"Path to the data file\")\n",
    "\n",
    "\n",
    "@register_function(config_type=DetectOutliersIQRConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])\n",
    "async def detect_outliers_iqr_function(config: DetectOutliersIQRConfig, _builder: Builder):\n",
    "    \"\"\"Detect outliers in sales data using the Interquartile Range (IQR) method.\"\"\"\n",
    "    import pandas as pd\n",
    "\n",
    "    df = pd.read_csv(config.data_path)\n",
    "\n",
    "    async def _detect_outliers_iqr(metric: str) -> str:\n",
    "        \"\"\"\n",
    "        Detect outliers in retail data using the IQR method.\n",
    "\n",
    "        Args:\n",
    "            metric: Specific metric to check for outliers\n",
    "\n",
    "        Returns:\n",
    "            Dictionary containing outlier analysis results\n",
    "        \"\"\"\n",
    "        if metric == \"None\":\n",
    "            column = \"Revenue\"\n",
    "        else:\n",
    "            column = metric\n",
    "\n",
    "        q1 = df[column].quantile(0.25)\n",
    "        q3 = df[column].quantile(0.75)\n",
    "        iqr = q3 - q1\n",
    "        outliers = df[(df[column] < q1 - 1.5 * iqr) | (df[column] > q3 + 1.5 * iqr)]\n",
    "\n",
    "        return f\"Outliers in {column} are {outliers.to_dict('records')}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(\n",
    "        _detect_outliers_iqr,\n",
    "        description=_detect_outliers_iqr.__doc__)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lkEBP-0O59Fd"
   },
   "source": [
    "<a id=\"registering-tools\"></a>\n",
    "## 2.4) Registering Tools\n",
    "\n",
    "We need to update the `register.py` file to register these tools with NeMo Agent toolkit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile -a retail_sales_agent_nb3/src/retail_sales_agent_nb3/register.py\n",
    "\n",
    "from . import sales_per_day_tool\n",
    "from . import detect_outliers_tool\n",
    "from . import total_product_sales_data_tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "adWCjEEb66H8"
   },
   "source": [
    "<a id=\"update-config\"></a>\n",
    "## 2.5) Updating The Configuration File\n",
    "\n",
    "Below we show how to update the default configuration file for this new tool-calling retail sales agent with the new tools (python methods) that have been defined and registered properly.\n",
    "\n",
    "Take a moment to analyze the new `retail_sales_agent_nb3/configs/config.yml` file below, where a `functions` header has been defined and each registered tool from the previous section is listed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb3/configs/config.yml\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.3-70b-instruct\n",
    "    temperature: 0.0\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "functions:\n",
    "  total_product_sales_data:\n",
    "    _type: get_total_product_sales_data\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  sales_per_day:\n",
    "    _type: get_sales_per_day\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  detect_outliers:\n",
    "    _type: detect_outliers_iqr\n",
    "    data_path: data/retail_sales_data.csv\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names:\n",
    "    - total_product_sales_data\n",
    "    - sales_per_day\n",
    "    - detect_outliers\n",
    "  llm_name: nim_llm\n",
    "  verbose: true\n",
    "  handle_parsing_errors: true\n",
    "  max_retries: 2\n",
    "  description: \"A helpful assistant that can answer questions about the retail sales CSV data\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_gkpzmGp7VaD"
   },
   "source": [
    "<a id=\"running-initial-workflow\"></a>\n",
    "## 2.6) Running the Initial Workflow\n",
    "\n",
    "The workflow has been properly created, new tools defined, registered, and incorporated into the config.yml. We are now ready to run the agent and test out that the ReAct agent is able to properly determine the intent of the input query and complete the necessary tool calling to serve the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "nat workflow reinstall retail_sales_agent_nb3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7vnTjDYzCYrs"
   },
   "source": [
    "This first query asks how laptop sales compare to phone sales.\n",
    "\n",
    "In the output, we expect to see calls to the `total_product_sales_data` tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file=retail_sales_agent_nb3/configs/config.yml --input \"How do laptop sales compare to phone sales?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y8SfmI35CNKn"
   },
   "source": [
    "In this next query we ask what were the laptop sales on a specific date.\n",
    "\n",
    "In the output, we expect to see a call to the `sales_per_day` tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file=retail_sales_agent_nb3/configs/config.yml \\\n",
    "    --input \"What were the laptop sales on February 16th 2024?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sk6-fqgjDE9s"
   },
   "source": [
    "In the last query we ask if there were any outliers in sales.\n",
    "\n",
    "In the output, we expect to see a call to the `detect_outliers` tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file=retail_sales_agent_nb3/configs/config.yml --input \"What were the outliers in 'Revenue'?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F_f7mZc0FEff"
   },
   "source": [
    "<a id=\"adding-retrieval-tool\"></a>\n",
    "# 3.0) Adding a Retrieval Tool using LlamaIndex\n",
    "\n",
    "After basic tool calling has been demonstrated, adding context retrieval tools to our agent is a reasonable next step.\n",
    "\n",
    "In this section we will equip our agent with a tool that is capable of performing retrieval of additional context to answer questions about some new consumer products that the backbone model likely doesn't have pretrained knowledge of. It will use a vector store that stores details about products. We can create this agent using LlamaIndex to demonstrate the framework-agnostic capability of the library."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S_WidV3lSV0E"
   },
   "source": [
    "<a id=\"defining-retrieval-tool\"></a>\n",
    "## 3.1) Defining the Retrieval Tool\n",
    "\n",
    "Just like with section 2 above, we will define our new tool by writing to a new source file for this agent: `retail_sales_rag_tool.py`. This tool using Llama Index to chunk, embed, index, and retrieve ranked results from the source text when called.\n",
    "\n",
    "<span style=\"color: red\"><i>Note: In a real‑world scenario, it is not recommended to upsert records at query time due to latency. However, the simplistic approach below is adequate for this demo.</i></span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb3/src/retail_sales_agent_nb3/retail_sales_rag_tool.py\n",
    "import logging\n",
    "import os\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from nat.builder.builder import Builder\n",
    "from nat.builder.framework_enum import LLMFrameworkEnum\n",
    "from nat.builder.function_info import FunctionInfo\n",
    "from nat.cli.register_workflow import register_function\n",
    "from nat.data_models.component_ref import EmbedderRef\n",
    "from nat.data_models.component_ref import LLMRef\n",
    "from nat.data_models.function import FunctionBaseConfig\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class LlamaIndexRAGConfig(FunctionBaseConfig, name=\"retail_sales_rag\"):\n",
    "\n",
    "    llm_name: LLMRef = Field(description=\"The name of the LLM to use for the RAG engine.\")\n",
    "    embedder_name: EmbedderRef = Field(description=\"The name of the embedder to use for the RAG engine.\")\n",
    "    data_dir: str = Field(description=\"The directory containing the data to use for the RAG engine.\")\n",
    "    description: str = Field(description=\"A description of the knowledge included in the RAG system.\")\n",
    "    collection_name: str = Field(default=\"context\", description=\"The name of the collection to use for the RAG engine.\")\n",
    "\n",
    "\n",
    "def _walk_directory(root: str):\n",
    "    for root, dirs, files in os.walk(root):\n",
    "        for file_name in files:\n",
    "            yield os.path.join(root, file_name)\n",
    "\n",
    "\n",
    "@register_function(config_type=LlamaIndexRAGConfig, framework_wrappers=[LLMFrameworkEnum.LLAMA_INDEX])\n",
    "async def retail_sales_rag_tool(config: LlamaIndexRAGConfig, builder: Builder):\n",
    "    from llama_index.core import Settings\n",
    "    from llama_index.core import SimpleDirectoryReader\n",
    "    from llama_index.core import StorageContext\n",
    "    from llama_index.core import VectorStoreIndex\n",
    "    from llama_index.core.node_parser import SentenceSplitter\n",
    "\n",
    "    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)\n",
    "    embedder = await builder.get_embedder(config.embedder_name, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)\n",
    "\n",
    "    Settings.embed_model = embedder\n",
    "    Settings.llm = llm\n",
    "\n",
    "    files = list(_walk_directory(config.data_dir))\n",
    "    docs = SimpleDirectoryReader(input_files=files).load_data()\n",
    "    logger.info(\"Loaded %s documents from %s\", len(docs), config.data_dir)\n",
    "\n",
    "    parser = SentenceSplitter(\n",
    "        chunk_size=400,\n",
    "        chunk_overlap=20,\n",
    "        separator=\" \",\n",
    "    )\n",
    "    nodes = parser.get_nodes_from_documents(docs)\n",
    "\n",
    "    index = VectorStoreIndex(nodes)\n",
    "\n",
    "    query_engine = index.as_query_engine(similarity_top_k=3, )\n",
    "\n",
    "    async def _arun(inputs: str) -> str:\n",
    "        \"\"\"\n",
    "        Search product catalog for information about tablets, laptops, and smartphones\n",
    "        Args:\n",
    "            inputs: user query about product specifications\n",
    "        \"\"\"\n",
    "        try:\n",
    "            response = query_engine.query(inputs)\n",
    "            return str(response.response)\n",
    "\n",
    "        except Exception as e:\n",
    "            logger.error(\"RAG query failed: %s\", e)\n",
    "            return f\"Sorry, I couldn't retrieve information about that product. Error: {str(e)}\"\n",
    "\n",
    "    yield FunctionInfo.from_fn(_arun, description=config.description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we will register it..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile -a retail_sales_agent_nb3/src/retail_sales_agent_nb3/register.py\n",
    "\n",
    "from . import retail_sales_rag_tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And configure our agent to use it..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qfhxCrUWOT1R"
   },
   "source": [
    "<a id=\"retrieval-config\"></a>\n",
    "## 3.2) Retrieval Tool Workflow Configuration File\n",
    "\n",
    "We need a new workflow configuration file which incorporates this new tool."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i3mqYNQxxVRu"
   },
   "source": [
    "The key additions are:\n",
    "* Introduction of an Embedder (`nvidia/nv-embedqa-e5-v5`)\n",
    "* Addition of an instantiated `retail_sales_rag` tool which processes files in the `data/rag` directory\n",
    "* A custom RAG agent which interfaces with the RAG tool, providing a natural language frontend to the tool.\n",
    "* Adding the custom RAG agent to the list of available tools to our original agent.\n",
    "\n",
    "> **Note:** _The only impactful change to the top-level agent was the addition of the new RAG agent. All other changes to the configuration were for enabling the RAG agent._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile retail_sales_agent_nb3/configs/config_rag.yml\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.3-70b-instruct\n",
    "    temperature: 0.0\n",
    "    max_tokens: 2048\n",
    "    context_window: 32768\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "embedders:\n",
    "  nim_embedder:\n",
    "    _type: nim\n",
    "    model_name: nvidia/nv-embedqa-e5-v5\n",
    "    truncate: END\n",
    "    api_key: $NVIDIA_API_KEY\n",
    "\n",
    "functions:\n",
    "  total_product_sales_data:\n",
    "    _type: get_total_product_sales_data\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  sales_per_day:\n",
    "    _type: get_sales_per_day\n",
    "    data_path: data/retail_sales_data.csv\n",
    "  detect_outliers:\n",
    "    _type: detect_outliers_iqr\n",
    "    data_path: data/retail_sales_data.csv\n",
    "\n",
    "  product_catalog_rag:\n",
    "    _type: retail_sales_rag\n",
    "    llm_name: nim_llm\n",
    "    embedder_name: nim_embedder\n",
    "    collection_name: product_catalog_rag\n",
    "    data_dir: data/rag/\n",
    "    description: \"Search product catalog for TabZen tablet, AeroBook laptop, NovaPhone specifications\"\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names:\n",
    "    - total_product_sales_data\n",
    "    - sales_per_day\n",
    "    - detect_outliers\n",
    "    - product_catalog_rag\n",
    "  llm_name: nim_llm\n",
    "  max_history: 10\n",
    "  max_iterations: 15\n",
    "  description: \"A helpful assistant that can answer questions about the retail sales CSV data\"\n",
    "  verbose: true"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "COig5flmR8f8"
   },
   "source": [
    "<a id=\"running-retrieval-workflow\"></a>\n",
    "## 3.3) Running the Workflow\n",
    "\n",
    "We can now test the RAG-enabled workflow with the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file=retail_sales_agent_nb3/configs/config_rag.yml \\\n",
    "    --input \"What is the Ark S12 Ultra tablet and what are its specifications?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the significance of what we've achieved in just a few lines of code: a reasoning agent was brought up with tool calls that allow it the execute predefined python functions to achieve what an LLM alone cannot. Additionally, we've incorporated context retrieval RAG into the same workflow so that the agent can access domain-specific or real time data sources that it's backbone LLM has never seen during training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"next-steps\"></a>\n",
    "# 4.0) Next Steps\n",
    "\n",
    "In the next notebook in our series: `mcp_setup_and_integration.ipynb` we will demonstrate how distributed sets of tools can be hosted on MCP servers and used by MCP clients."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "unew_312",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
